package top.ezzz.ai;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import lombok.extern.slf4j.Slf4j;
import org.noear.solon.ai.chat.ChatModel;
import org.noear.solon.ai.chat.ChatResponse;
import org.noear.solon.ai.chat.ChatSession;
import org.noear.solon.ai.chat.message.ChatMessage;
import org.noear.solon.ai.embedding.EmbeddingModel;
import org.noear.solon.ai.rag.Document;
import org.noear.solon.ai.rag.Repository;
import org.noear.solon.ai.rag.loader.TextLoader;
import org.noear.solon.ai.rag.repository.InMemoryRepository;
import org.noear.solon.ai.rag.splitter.RegexTextSplitter;
import org.noear.solon.ai.rag.splitter.SplitterPipeline;
import org.noear.solon.ai.rag.splitter.TokenSizeTextSplitter;
import org.noear.solon.ai.rag.util.QueryCondition;
import org.noear.solon.annotation.*;
import org.noear.solon.core.handle.Context;
import org.noear.solon.core.handle.UploadedFile;
import org.noear.solon.core.util.MimeType;
import org.noear.solon.web.cors.annotation.CrossOrigin;
import org.noear.solon.web.sse.SseEvent;
import reactor.core.publisher.Flux;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import static top.ezzz.ai.SolonAiSession.chatSessionMap;

@Slf4j
@CrossOrigin(origins = "*")
@Controller
public class SolonAiController {

    @Inject("qwenChatModel")
    private ChatModel chatModel;

    @Inject("bochaWebRepository")
    private Repository webRepository;

    @Inject("inMemoryRepository")
    private InMemoryRepository memoryRepository;

    @Inject("bgeEmbeddingModel")
    private EmbeddingModel embeddingModel;

    //简单返回
    @Mapping("chat")
    public ChatMessage chat(@Body() String message) throws IOException {
        return chatModel.prompt(message).call().getMessage();
    }


    @Produces(MimeType.TEXT_EVENT_STREAM_UTF8_VALUE)
    //@Produces(MimeType.APPLICATION_X_NDJSON_UTF8_VALUE)
    @Mapping("chat-sse")
    public Flux<SseEvent> chatSse(String message,Context ctx) throws IOException {
        log.info("message = {}", message);
        //web检索
        //List<Document> context = webRepository.search(new QueryCondition(message).limit(4));
        //log.info("webRepository context = {}", JSONUtil.toJsonStr(context));
        //内存检索
        //List<Document> context = memoryRepository.search(new QueryCondition(message).limit(4));
        //log.info("memoryRepository context = {}", JSONUtil.toJsonStr(context));
        //消息增强
        //ChatMessage chatMessage = ChatMessage.augment(message, null);
        //保留上下文
        String sessionId = ctx.sessionId();
        ChatSession chatSession = chatSessionMap.get(sessionId);
        if (chatSession == null) {
            chatSession = new ChatSessionLimit(sessionId, 10);
            chatSessionMap.put(sessionId, chatSession);
        }
        chatSession.addMessage(ChatMessage.ofUser(message));

        // 调用模型
        return Flux.from(chatModel.prompt(chatSession).stream())
                .doOnError(error -> log.error(error.getMessage(), error))
                .doOnComplete(() -> log.info("chatModel stream complete"))
                .map(ChatResponse::getMessage)
//                .map(msg -> new SseEvent().data(StrUtil.nullToEmpty(msg.getContent())));
                .map(msg -> new SseEvent().data(JSONUtil.parseObj(msg)));
    }

    @Mapping("chat-session")
    public ChatMessage chatSession(Context ctx, String message) throws IOException {
        String sessionId = ctx.sessionId();
        ChatSession chatSession = chatSessionMap.get(sessionId);
        if (chatSession == null) {
            chatSession = new ChatSessionLimit(sessionId, 10);
            chatSessionMap.put(sessionId, chatSession);
        }
        chatSession.addMessage(ChatMessage.ofUser(message));
        log.info("size = {}", chatSession.getMessages().size());
        ChatMessage resultMsg = chatModel.prompt(chatSession).call().getMessage();
        chatSession.addMessage(resultMsg);
        return resultMsg;
    }


    @Mapping("upload")
    public boolean upload(UploadedFile file) throws IOException {
        // 创建文件
        File uploadFile = FileUtil.touch("src/main/resources/chat.txt");
        // 保存文件
        file.transferTo(uploadFile);
        //加载器
        List<Document> loaded = new TextLoader(uploadFile).load();
        //分割器，再次（按需）组合切割
        List<Document> documents = new SplitterPipeline()
                .next(new RegexTextSplitter("\r\n\r\n"))
                .next(new TokenSizeTextSplitter(500))
                .split(loaded);

        int batch = documents.size() / 10;
        //batch = 10;
        for (int index = 0; index <= batch; index++) {
            List<Document> subList = CollUtil.sub(documents, index * 10, (index + 1) * 10);
            log.info("batch = {} / {} , subList = {}", index, batch, subList.size());
            //存储仓库, 会自动进行embedding
            memoryRepository.insert(subList);
        }
        return true;
    }


}